import cv2
import yaml
from detector import Detector
from tracker import Sort
from counter import Counter
import os
import logging
from collections import defaultdict
import numpy as np

# 1. 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def main(config_path):
    """
    主函数，运行整个处理流程。
    :param config_path: 配置文件的路径。
    """
    # 1. 加载配置
    logging.info(f"Loading config from: {config_path}")
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)

    video_source = config['video_source']
    model_path = config['yolo_model_path']
    line_points = config['line_points']

    # 2. 初始化模块
    logging.info("Initializing Detector, Tracker, and Counter modules.")
    detector = Detector(model_path)
    tracker = Sort(max_age=20, min_hits=3, iou_threshold=0.3)
    counter = Counter(line_points)

    # 3. 打开视频源
    logging.info(f"Opening video source: {video_source}")
    cap = cv2.VideoCapture(video_source)
    if not cap.isOpened():
        logging.error(f"Could not open video source: {video_source}")
        return

    # 获取视频的宽度、高度和帧率
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    logging.info(f"Video properties: {width}x{height} @ {fps} FPS")

    # 4. 初始化视频写入器
    output_path = 'output.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    logging.info(f"Output video will be saved to: {output_path}")

    # 用于存储轨迹
    trajectories = defaultdict(list)

    frame_count = 0
    # 5. 逐帧处理
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            logging.info("End of video stream.")
            break
        
        frame_count += 1
        if frame_count % 100 == 0:
            logging.info(f"Processing frame {frame_count}")

        # a. 目标检测
        detections = detector.detect(frame)

        # b. 目标跟踪
        if detections.size > 0:
            tracks = tracker.update(detections)
        else:
            # 如果没有检测到目标，仍然需要调用tracker.update()来更新卡尔曼滤波器的状态
            tracks = tracker.update()

        # c. 计数
        in_ids, out_ids = [], []
        if len(tracks) > 0:
            in_ids, out_ids = counter.update(tracks)
        
        if in_ids:
            logging.info(f"Person IN: {in_ids}")
        if out_ids:
            logging.info(f"Person OUT: {out_ids}")

        # d. 可视化
        # 绘制计数线
        cv2.line(frame, tuple(line_points[0]), tuple(line_points[1]), (0, 255, 0), 2)

        # 绘制跟踪结果和轨迹
        for track in tracks:
            x1, y1, x2, y2, track_id = track.astype(int)
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            
            # 更新轨迹
            trajectories[track_id].append((center_x, center_y))
            
            # 绘制轨迹
            if len(trajectories[track_id]) > 1:
                cv2.polylines(frame, [np.array(trajectories[track_id], dtype=np.int32)], isClosed=False, color=(255, 255, 0), thickness=2)

            # 绘制跟踪框
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(frame, f"ID: {track_id}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

        # 绘制计数结果
        cv2.putText(frame, f"In: {counter.in_counts}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        cv2.putText(frame, f"Out: {counter.out_counts}", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        # e. 写入输出视频
        out.write(frame)

        # f. 显示结果 (可选，在非GUI环境中可以注释掉)
        # cv2.imshow('Bus Passenger Counting', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 6. 释放资源
    logging.info("Releasing resources.")
    cap.release()
    out.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # Get the absolute path to the directory of the current script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    # Construct the absolute path to the config file
    config_path = os.path.join(script_dir, '..', 'configs', 'config.yaml')
    # Normalize the path to resolve any '..' components
    config_path = os.path.normpath(config_path)
    main(config_path)